from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, SubmitField, BooleanField, IntegerField, SelectField
from wtforms.validators import DataRequired, Length, Email, EqualTo, ValidationError
from app.models import User, Student

class LoginForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired(), Length(min=2, max=20)])
    password = PasswordField('密码', validators=[DataRequired()])
    remember = BooleanField('记住我')
    submit = SubmitField('登录')

class RegistrationForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired(), Length(min=2, max=20)])
    email = StringField('邮箱', validators=[DataRequired(), Email()])
    password = PasswordField('密码', validators=[DataRequired(), Length(min=6)])
    confirm_password = PasswordField('确认密码', validators=[DataRequired(), EqualTo('password')])
    submit = SubmitField('注册')

    def validate_username(self, username):
        user = User.query.filter_by(username=username.data).first()
        if user:
            raise ValidationError('该用户名已存在，请选择其他用户名。')

    def validate_email(self, email):
        user = User.query.filter_by(email=email.data).first()
        if user:
            raise ValidationError('该邮箱已被注册，请使用其他邮箱。')

class StudentForm(FlaskForm):
    student_id = StringField('学号', validators=[DataRequired(), Length(min=4, max=4)])
    name = StringField('姓名', validators=[DataRequired(), Length(min=2, max=50)])
    gender = SelectField('性别', choices=[('男', '男'), ('女', '女')], validators=[DataRequired()])
    age = IntegerField('年龄', validators=[DataRequired()])
    major_id = SelectField('专业', coerce=int, validators=[DataRequired()])
    submit = SubmitField('提交')

    def validate_student_id(self, student_id):
        # 检查是否是4位数字
        if not student_id.data.isdigit():
            raise ValidationError('学号必须为4位数字。')
        # 检查是否是编辑模式，如果是则跳过检查
        if hasattr(self, 'student_id_initial') and self.student_id_initial == student_id.data:
            return
        student = Student.query.filter_by(student_id=student_id.data).first()
        if student:
            raise ValidationError('该学号已存在。')